package main

import (
	"sync"
	"time"
)

// 实现一个任务池
// 1. 添加一个任务到队列中
// 2. 取消一个还在队列中的任务
// 3. 任务不可以重复添加
func main() {
	taskPool := NewTaskPool()
	taskPool.Run()

	// 示例：添加任务
	task1 := &Task{
		Id: "1", Name: "Task1",
		runTask: func() {
			time.Sleep(1 * time.Second)
			print("Task1 executed\n")
		},
		onCancel: func() {
			print("Task1 cancelled\n")
		},
	}
	taskPool.AddTask(task1)

	// 示例：修改并发任务数量
	taskPool.SetConcurrentCount(3)
}

type Task struct {
	Id       string
	Name     string
	runTask  func()
	onCancel func() // 取消任务时执行的方法
}

type TaskPool struct {
	tasks []*Task
	sync.Mutex
	semaphore  chan struct{} // 信号量，用于控制并发任务数量
	concurrent int           // 当前并发任务数量
}

func NewTaskPool() *TaskPool {
	return &TaskPool{
		tasks:      make([]*Task, 0),
		semaphore:  make(chan struct{}, 5), // 初始并发数为 5
		concurrent: 5,
	}
}

// AddTask 实现添加一个任务到队列中
func (tp *TaskPool) AddTask(t *Task) {
	tp.Mutex.Lock()
	defer tp.Mutex.Unlock()
	// 检查任务是否已经存在
	for _, task := range tp.tasks {
		if task.Id == t.Id {
			return
		}
	}
	tp.tasks = append(tp.tasks, t)
}

// CancelTask 实现取消一个还在队列中的任务
func (tp *TaskPool) CancelTask(t *Task) {
	tp.Mutex.Lock()
	defer tp.Mutex.Unlock()
	for i, task := range tp.tasks {
		if task.Id == t.Id {
			if task.onCancel != nil {
				task.onCancel() // 执行取消任务时的特定方法
			}
			tp.tasks = append(tp.tasks[:i], tp.tasks[i+1:]...)
			break
		}
	}
}

// Run 执行任务
func (tp *TaskPool) Run() {
	go func() {
		for {
			tp.Mutex.Lock()
			if len(tp.tasks) > 0 {
				for i := 0; i < len(tp.tasks); i++ {
					task := tp.tasks[i]
					tp.semaphore <- struct{}{} // 获取信号量
					go func(t *Task) {
						defer func() { <-tp.semaphore }() // 释放信号量
						t.runTask()
					}(task)
				}
			}
			tp.Mutex.Unlock()
			time.Sleep(time.Microsecond * 100)
		}
	}()
}

// SetConcurrentCount 设置并发任务数量
func (tp *TaskPool) SetConcurrentCount(count int) {
	tp.Mutex.Lock()
	defer tp.Mutex.Unlock()
	tp.concurrent = count
	newSemaphore := make(chan struct{}, count)
	// 调整信号量容量，释放多余的信号量
	for i := 0; i < count; i++ {
		if i < len(tp.semaphore) {
			newSemaphore <- struct{}{}
		}
	}
	tp.semaphore = newSemaphore
}
